#!/usr/bin/python

""" Measure cpu usage, disk io, network io of cluster
argv[1]: result_file_name

"""

import os, subprocess, time, sys
import re
import socket

def digit_2_str(value):
    pattern = re.compile(r'(?<=\d)(?=(\d\d\d)+(?!\d))')
    return pattern.sub(',', str(value))

def get_curr_time_str():
    return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())

def get_network_bytes():
    recv_bytes = 0;
    send_bytes = 0;
    
    f = open('/proc/net/dev')
    lines = f.read().split('\n')
    
    for line in lines:
        curr_line = line.strip()
        
        if curr_line.startswith('eth'):
            i = curr_line.find(':')
            curr_line = curr_line[i + 1 : ]
            items = curr_line.split()
            recv_bytes = int(items[0])
            send_bytes = int(items[8])
            
    f.close()
    
    return recv_bytes, send_bytes


def get_cpu_jiffies():
    idle_jiffies = 0
    total_jiffies = 0
    f = open('/proc/stat')
    line = f.read().split('\n')[0]
    items = line.split()
    
    for item in items[1:]:
        total_jiffies += int(item)
    
    idle_jiffies = int(items[4])
    
    f.close()
    return idle_jiffies, total_jiffies

def get_disk_io():
    read_sector_count = 0
    read_time = 0
    write_sector_count = 0
    write_time = 0

    f = open('/proc/diskstats')
    lines = f.read().split('\n')
    for line in lines:
        if 'xvda1' in line:
            items = line.split()
            read_sector_count = int(items[5])
            read_time = int(items[6])
            write_sector_count = int(items[9])
            write_time = int(items[10])
    f.close()
    return read_sector_count, read_time, write_sector_count, write_time
            
def printf(filename, info):
    f = open(filename, 'a')
    f.write(info)
    f.flush()
    f.close()


if __name__ == '__main__':
    pig_command = '/root/software/apache-cassandra-0.7.4-src/contrib/pig/bin/pig_cassandra_measure'
    pig_script = '/root/pigappscale/shopping_topCustomer.pig'
    command = [pig_command, pig_script]
    #command = ['/root/test_code/pig_bear/measure1/bear_pig_cassandra', '-x', 'local', '/root/test_code/pig_bear/measure_cluster/guestbook.pig']
    #command = ['/root/test_code/pig_bear/measure1/bear_pig_cassandra', '/root/test_code/pig_bear/measure_cluster/guestbook.pig']            
    master_ip = '128.111.55.223'
    slave_ip_list = ['128.111.55.204', '128.111.55.208']
    start_msg_str = 'BEAR_START'
    start_ok_msg_str = 'BEAR_START_OK'
    end_msg_str = 'BEAR_END'
     
    slave_port = 30000
    
    # get local ip
    s = os.popen('ifconfig eth0 | grep "inet "').read()
    i = s.find(':')
    s = s[i + 1 :]
    local_ip = s.split()[0]
    
    print 'Local IP : ' + local_ip
    
    if master_ip != local_ip:
        # is slave
        print "I'm slave"
        print get_curr_time_str() + ' : Listening on port ' + str(slave_port) + ' ... '
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.bind((local_ip , slave_port))
        sock.listen(5)
        
        while True:
            cs, address  =  sock.accept()
            print get_curr_time_str() + ' : Accepted, Client ip is : ' + str(address)
            s = cs.recv(1024)
            print get_curr_time_str() + ' : Received string : [' + s + ']'
            if s == start_msg_str:
                # record system informations before running
                time_old = time.time()
                recv_bytes_old, send_bytes_old = get_network_bytes()
                idle_jiffies_old, total_jiffies_old = get_cpu_jiffies()
                read_sector_count_old, read_time_old, write_sector_count_old, write_time_old = get_disk_io()
                
                print get_curr_time_str() + ' : Send ' + start_ok_msg_str + ' to ' + str(address)
                cs.send(start_ok_msg_str)
                
                print get_curr_time_str() + ' : Waiting ... '
                s = cs.recv(1024)
                print get_curr_time_str() + ' : Receive string : [' + s + ']'
                
                if s == end_msg_str:
                    # record system informations after running
                    time_curr = time.time()
                    recv_bytes_curr, send_bytes_curr = get_network_bytes()
                    idle_jiffies_curr, total_jiffies_curr = get_cpu_jiffies()
                    read_sector_count_curr, read_time_curr, write_sector_count_curr, write_time_curr = get_disk_io()
                    
                    
                    # calculate results
                    print get_curr_time_str() + ' : Calculating reuslts ... '
                    total_seconds = time_curr - time_old
                    total_network_recv = recv_bytes_curr - recv_bytes_old
                    total_network_send = send_bytes_curr - send_bytes_old
                    total_cpu_jiffies = total_jiffies_curr - total_jiffies_old
                    total_cpu_idle_jiffies = idle_jiffies_curr - idle_jiffies_old
                    total_disk_read_sector = read_sector_count_curr - read_sector_count_old
                    total_disk_read_time = read_time_curr - read_time_old
                    total_disk_write_sector = write_sector_count_curr - write_sector_count_old
                    total_disk_write_time = write_time_curr - write_time_old
                    
                    result_str = str(total_seconds) + ', ' + str(total_cpu_jiffies) + ', ' + str(total_cpu_idle_jiffies) + ', ' \
                                 + str(total_disk_read_sector) + ', ' + str(total_disk_read_time) + ', ' + str(total_disk_write_sector) + ', ' \
                                 + str(total_disk_write_time) + ', ' + str(total_network_recv) + ', ' + str(total_network_send)
                    
                    print 'Sending result "' + result_str + '" to ' + str(address)
                    cs.send(result_str)
                    print 'done!\n'
        
    else:
        # is master
        print "I'm master"
        
        sock_list = []
        for slave_ip in slave_ip_list:
            # connect to slave
            sock = socket.socket()
            print get_curr_time_str() + ' : Connecting to ' + slave_ip + ' : ' + str(slave_port)
            sock.connect((slave_ip , slave_port))
            
            print get_curr_time_str() + ' : Sending ' + start_msg_str + ' to slave (' + slave_ip + ')'
            sock.send(start_msg_str)
            print get_curr_time_str() + ' : Waiting for ack'
            s = sock.recv(1024)
            if s == start_ok_msg_str:
                sock_list.append((slave_ip, sock))
        
        if len(sock_list) == len(slave_ip_list):
            # record system informations before running
            time_old = time.time()
            recv_bytes_old, send_bytes_old = get_network_bytes()
            idle_jiffies_old, total_jiffies_old = get_cpu_jiffies()
            read_sector_count_old, read_time_old, write_sector_count_old, write_time_old = get_disk_io()
            
            # run pig_cassandra
            print get_curr_time_str() + ' : Running ... ' 
            p = subprocess.Popen(command)
            p.wait()
            #time.sleep(10);
            
            # record system informations after running
            time_curr = time.time()
            recv_bytes_curr, send_bytes_curr = get_network_bytes()
            idle_jiffies_curr, total_jiffies_curr = get_cpu_jiffies()
            read_sector_count_curr, read_time_curr, write_sector_count_curr, write_time_curr = get_disk_io()
            
            # get informations from slaves
            slave_result_dict = {}
            for sock_items in sock_list:
                slave_ip = sock_items[0]
                sock = sock_items[1]
                print get_curr_time_str() + ' : Getting results from ' + slave_ip
                sock.send(end_msg_str)
                
                s = sock.recv(1024)
                items = s.split(',')
                slave_result_dict[slave_ip] = items
                sock.close()
            
            # calculate results of master
            total_seconds = time_curr - time_old
            total_network_recv = recv_bytes_curr - recv_bytes_old
            total_network_send = send_bytes_curr - send_bytes_old
            total_cpu_jiffies = total_jiffies_curr - total_jiffies_old
            total_cpu_idle_jiffies = idle_jiffies_curr - idle_jiffies_old
            total_disk_read_sector = read_sector_count_curr - read_sector_count_old
            total_disk_read_time = read_time_curr - read_time_old
            total_disk_write_sector = write_sector_count_curr - write_sector_count_old
            total_disk_write_time = write_time_curr - write_time_old
            
            # display results
            first_col_width = 18
            col_width = 21
            
            
            header_str = ''.rjust(first_col_width) + master_ip.rjust(col_width) 
            for slave_ip in slave_ip_list:
                header_str += slave_ip.rjust(col_width)
	    header_str += 'Total'.rjust(col_width)
	    result_str = header_str
            
            result_str += '\nTotal Time'.ljust(first_col_width) + (str(total_seconds) + ' s').rjust(col_width)
            for slave_ip in slave_ip_list:
                result_str += (slave_result_dict[slave_ip][0] + ' s').rjust(col_width)
	    result_str += (str(total_seconds) + ' s').rjust(col_width)            

            cpu_jiffies_used = total_cpu_jiffies - total_cpu_idle_jiffies
            cpu_percent = round(cpu_jiffies_used * 100.0 / total_cpu_jiffies, 2)
            result_str += '\n' + 'CPU Usage'.ljust(first_col_width) + (digit_2_str(cpu_jiffies_used)  + ' jif (' + str(cpu_percent) + ' %)').rjust(col_width)
            for slave_ip in slave_ip_list:
                cpu_jiffies_total_slave = int(slave_result_dict[slave_ip][1])
                cpu_jiffies_idle_slave = int(slave_result_dict[slave_ip][2])
                cpu_jiffies_used_slave = cpu_jiffies_total_slave - cpu_jiffies_idle_slave
                cpu_percent_slave = round(cpu_jiffies_used_slave * 100.0 / cpu_jiffies_total_slave, 2)
                result_str += (digit_2_str(cpu_jiffies_used_slave)  + ' jif (' + str(cpu_percent_slave) + ' %)').rjust(col_width)
		
		cpu_jiffies_used += cpu_jiffies_used_slave
		total_cpu_jiffies += cpu_jiffies_total_slave
	    cpu_percent = round(cpu_jiffies_used * 100.0 / total_cpu_jiffies, 2)
	    result_str += (digit_2_str(cpu_jiffies_used)  + ' jif (' + str(cpu_percent) + ' %)').rjust(col_width)

            
            result_str += '\n' + 'Disk Read'.ljust(first_col_width) + (digit_2_str(total_disk_read_sector) + ' sectors').rjust(col_width)
            for slave_ip in slave_ip_list:
                result_str += (digit_2_str(int(slave_result_dict[slave_ip][3])) + ' sectors').rjust(col_width)
		total_disk_read_sector += int(slave_result_dict[slave_ip][3])
	    result_str += (digit_2_str(total_disk_read_sector) + ' sectors').rjust(col_width)
            
            result_str += '\n' + 'Disk Read time'.ljust(first_col_width) + (digit_2_str(total_disk_read_time) + ' ms').rjust(col_width)
            for slave_ip in slave_ip_list:
                result_str += (digit_2_str(int(slave_result_dict[slave_ip][4])) + ' ms').rjust(col_width)
		total_disk_read_time += int(slave_result_dict[slave_ip][4])
	    result_str += (digit_2_str(total_disk_read_time) + ' ms').rjust(col_width)

            
            result_str += '\n' + 'Disk Write'.ljust(first_col_width) + (digit_2_str(total_disk_write_sector) + ' sectors').rjust(col_width)
            for slave_ip in slave_ip_list:
                result_str += (digit_2_str(int(slave_result_dict[slave_ip][5])) + ' sectors').rjust(col_width)
                total_disk_write_sector += int(slave_result_dict[slave_ip][5])
	    result_str += (digit_2_str(total_disk_write_sector) + ' sectors').rjust(col_width)
            

            result_str += '\n' + 'Disk Write time'.ljust(first_col_width) + (digit_2_str(total_disk_write_time) + ' ms').rjust(col_width)
            for slave_ip in slave_ip_list:
                result_str += (digit_2_str(int(slave_result_dict[slave_ip][6])) + ' ms').rjust(col_width)
		total_disk_write_time += int(slave_result_dict[slave_ip][6])
	    result_str += (digit_2_str(total_disk_write_time) + ' ms').rjust(col_width)            
            

            result_str += '\n' + 'Network Recv'.ljust(first_col_width) + (digit_2_str(total_network_recv) + ' Bytes').rjust(col_width)
	    total_network_recv_all = total_network_recv
            for slave_ip in slave_ip_list:
                result_str += (digit_2_str(int(slave_result_dict[slave_ip][7])) + ' Bytes').rjust(col_width)
		total_network_recv_all += int(slave_result_dict[slave_ip][7])
	    result_str += (digit_2_str(total_network_recv_all) + ' Bytes').rjust(col_width)                
                

            result_str += '\n' + 'Network Recv Spd'.ljust(first_col_width) + (digit_2_str(int(round(total_network_recv / total_seconds))) + ' B/s').rjust(col_width)
            for slave_ip in slave_ip_list:
                total_sec = float(slave_result_dict[slave_ip][0])
                bytes_num = int(slave_result_dict[slave_ip][7])
                result_str += (digit_2_str(int(round(bytes_num / total_sec))) + ' B/s').rjust(col_width)
            result_str += (digit_2_str(int(round(total_network_recv_all / total_seconds))) + ' B/s').rjust(col_width)
            
            result_str += '\n' + 'Network Send'.ljust(first_col_width) + (digit_2_str(total_network_send) + ' Bytes').rjust(col_width)
	    total_network_send_all = total_network_send
            for slave_ip in slave_ip_list:
                result_str += (digit_2_str(int(slave_result_dict[slave_ip][8])) + ' Bytes').rjust(col_width)
                total_network_send_all += int(slave_result_dict[slave_ip][8])
	    result_str += (digit_2_str(total_network_send_all) + ' Bytes').rjust(col_width)	    

                
            result_str += '\n' + 'Network Send Spd'.ljust(first_col_width) + (digit_2_str(int(round(total_network_send / total_seconds))) + ' B/s').rjust(col_width)
            for slave_ip in slave_ip_list:
                total_sec = float(slave_result_dict[slave_ip][0])
                bytes_num = int(slave_result_dict[slave_ip][8])
                result_str += (digit_2_str(int(round(bytes_num / total_sec))) + ' B/s').rjust(col_width)
            result_str += (digit_2_str(int(round(total_network_send_all / total_seconds))) + ' B/s').rjust(col_width)
            
            # TODO: we can write result_str into file
           
	    print 'Command :   ' + ' '.join(command)
	    print '============================================================================================='
	    lines = result_str.split('\n')
            print '\033[45m' + lines[0] + '\033[0m'
            
            i = 0
            for line in lines[1: ]:
                i += 1
                if i % 2 == 0:
                    print '\033[46m' + line + '\033[0m'
                else:
                    print line
            
            print '============================================================================================='
	   
	    printf(sys.argv[1], result_str + '\n\n')
